import os
import glob
import logging
import glob
print(111)
from Crypto.Cipher import AES
print(111)
logging.basicConfig(level=logging.INFO,
                format='%(asctime)s [%(filename)s-%(lineno)d-%(levelname)s] %(message)s',
                datefmt='%Y-%m-%d %H:%M:%S',
                )


def readfile(finepath, mod="rb"):
    with open(finepath, mod) as fp:
        return fp.read()

def decrypt(data, key):
    try:
        aes = AES.new(key, AES.MODE_CBC, key)
        des = aes.decrypt(data)
        return des
    except ValueError:
        return b""


def decrypt_ts(src, key=None):
    with open(src, "rb") as fp:
        if not key:
            return fp.read()
        return decrypt(fp.read(), key)


def list_st(path):
    def sort(i):
        return int(i.split(".")[0])
    tss = []
    key = ""
    for i in os.listdir(path):
        if i.endswith(".ts"):
            tss.append(i)
        if i.endswith(".key"):
            key = open(os.path.join(path, i), "rb").read()
    tss.sort(key=sort)
    print(key)
    return tss, key
    


def ts_to_mp4(file_list, dst_name, key=None):
    os.makedirs(os.path.dirname(dst_name), exist_ok=True)
    with open(dst_name, "wb") as mp4_fp:
        for tsname in file_list:
            data = decrypt_ts(tsname, key)
            mp4_fp.write(data)

class TsToMp4():
    def __init__(self, tspath=None, mp4path=None):
        self.ts_path = tspath
        self.mp4path = mp4path
        self.key = None
        self.m3u8_list = []

    def get_m3u8(self):
        m3u8 = glob.glob(os.path.join(self.ts_path, "*.m3u8"))
        if not m3u8:
            logging.warning('%s m3u8 is not exists', self.ts_path)
            return []
        m3u8 = m3u8[0]
        with open(m3u8) as fp:
            return [line.strip() for line in fp.readlines() if line.strip()]

    def get_key(self, line):
        keyfile = line.split("URI=")[-1].split(",")[0].strip('"')
        return readfile(os.path.join(self.ts_path, keyfile))

    def main(self):
        m3u8_content_list = self.get_m3u8()
        for line in m3u8_content_list:
            if "EXT-X-KEY" in line:
                self.key = self.get_key(line)
            elif line.endswith(".ts"):
                ts_path = os.path.join(self.ts_path, line)
                if not os.path.exists(ts_path):
                    logging.warning("%s is not exists", ts_path)
                    continue
                self.m3u8_list.append(ts_path)
        ts_to_mp4(self.m3u8_list, self.mp4path, key=self.key)


for name in os.listdir("d:\\666"):
    path = os.path.join("d:\\666", name)
    print(path)
    TsToMp4(path, r"d:\\123\%s.mp4" % name).main()

